// Copyright 2020 The Tilt Brush Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

using UnityEngine;

namespace TiltBrush
{

    class ThickGeometryBrush : GeometryBrush
    {
        const int kVertsInSolid = 12;
        const int kSharedVertsInSolidPair = 6;
        const int kTrisInSolid = 8;
        const int kSharedTrisInSolidPair = 0;

        // Cross section is a diamond. Top and bottom
        // both have 3 verts (left, middle, right)
        //
        //       M
        //       .
        //      / \
        //  L ./___\. R
        //    ._____.
        //     \   /
        //      \ /
        //       '
        // First letter is Front/Back
        // Second letter is Left/Middle/Right
        // Third letter is Top/Bottom
        protected const int BRT = 0; // back  right
        protected const int BRB = 1; // back  right
        protected const int BMT = 2; // back  middle
        protected const int BMB = 3; // back  middle
        protected const int BLT = 4; // back  left
        protected const int BLB = 5; // back  left

        protected const int FRT = 6;  // front right
        protected const int FRB = 7;  // front right
        protected const int FMT = 8;  // front middle
        protected const int FMB = 9;  // front middle
        protected const int FLT = 10; // front left
        protected const int FLB = 11; // front left

        protected enum UVStyle
        {
            Distance,
            Stretch
        };

        const float kSolidMinLengthMeters_PS = 0.002f;
        const float kMinMoveLengthMeters_PS = 5e-4f;
        const float kBreakAngleScalar = 2.0f;
        const float kSolidAspectRatio = 0.2f;

        /// Amount of texture to chop off left and right edges, because
        /// interesting textures have ragged edges which don't work well when
        /// wrapped around tubes.
        [SerializeField] float m_TextureEdgeChop = .1f;

        [SerializeField]
        protected UVStyle m_uvStyle = UVStyle.Distance;

        public ThickGeometryBrush()
            : base(bCanBatch: true,
                upperBoundVertsPerKnot: kVertsInSolid,
                bDoubleSided: false)
        {
        }

        //
        // GeometryBrush API
        //

        protected override void InitBrush(BrushDescriptor desc, TrTransform localPointerXf)
        {
            base.InitBrush(desc, localPointerXf);
            m_geometry.Layout = GetVertexLayout(desc);
        }

        override public GeometryPool.VertexLayout GetVertexLayout(BrushDescriptor desc)
        {
            return new GeometryPool.VertexLayout
            {
                uv0Size = 2,
                uv1Size = 0,
                bUseNormals = true,
                bUseColors = true,
                bUseTangents = true,
            };
        }

        override public float GetSpawnInterval(float pressure01)
        {
            return kSolidMinLengthMeters_PS * App.METERS_TO_UNITS * POINTER_TO_LOCAL +
                (PressuredSize(pressure01) * kSolidAspectRatio);
        }

        override protected void ControlPointsChanged(int iKnot0)
        {
            // Frames knots, determines how much geometry each knot should get
            OnChanged_FrameKnots(iKnot0);

            ResizeGeometry();

            // Updating a control point affects geometry generated by previous knot
            // (if there is any). The HasGeometry check is not a micro-optimization:
            // it also keeps us from backing up past knot 0.
            {
                int start = (m_knots[iKnot0 - 1].HasGeometry) ? iKnot0 - 1 : iKnot0;
                OnChanged_MakeVertsAndNormals(start);
            }

            // Distance between knots isn't changed by the geometry, so there's not (yet?)
            // any reason to back up when creating them.
            if (m_uvStyle == UVStyle.Stretch)
            {
                OnChanged_StretchUVs(iKnot0);
            }
            else
            {
                OnChanged_DistanceUVs(iKnot0);
            }

            {
                int start = (m_knots[iKnot0 - 1].HasGeometry) ? iKnot0 - 1 : iKnot0;
                OnChanged_Tangents(start);
            }
        }

        // Fills in any knot data needed for geometry generation.
        // - fill in length, nRight, nSurface, iVert, iTri
        // - calculate strip-break points
        void OnChanged_FrameKnots(int iKnot0)
        {
            float minMove = kMinMoveLengthMeters_PS * App.METERS_TO_UNITS * POINTER_TO_LOCAL;

            Knot prev = m_knots[iKnot0 - 1];
            for (int iKnot = iKnot0; iKnot < m_knots.Count; ++iKnot)
            {
                Knot cur = m_knots[iKnot];

                bool shouldBreak = false;

                Vector3 vMove = cur.point.m_Pos - prev.point.m_Pos;
                cur.length = vMove.magnitude;

                // Rather than use unstable math, just bail and don't change the geometry.
                if (cur.length < minMove)
                {
                    shouldBreak = true;
                }

                // invariant: nSurface = nMove x nRight
                // If single-sided, always point the frontside towards the brush. Causes twisting.
                Vector3 nMove = vMove / cur.length;
                Vector3 vPreferredRight = m_Desc.m_BackIsInvisible
                    ? Vector3.Cross(cur.point.m_Orient * Vector3.forward, nMove)
                    : prev.nRight;
                ComputeSurfaceFrameNew(
                    vPreferredRight, nMove, cur.point.m_Orient,
                    out cur.nRight, out cur.nSurface);

                // More break checking; replicates previous logic
                if (prev.HasGeometry)
                {
                    float fWidthHeightRatio = cur.length / PressuredSize(cur.smoothedPressure);
                    float fBreakAngle = Mathf.Atan(fWidthHeightRatio) * Mathf.Rad2Deg * kBreakAngleScalar;
                    Vector3 vPrevMove = prev.point.m_Pos - m_knots[iKnot - 2].point.m_Pos;
                    if (Vector3.Angle(vPrevMove, vMove) > fBreakAngle)
                    {
                        shouldBreak = true;
                    }
                }

                if (shouldBreak)
                {
                    // End of strip
                    cur.iTri = prev.iTri + prev.nTri;
                    cur.nTri = 0;
                    cur.iVert = (ushort)(prev.iVert + prev.nVert);
                    cur.nVert = 0;
                    cur.nRight = cur.nSurface = Vector3.zero;
                }
                else
                {
                    // Beginning or middle of strip
                    cur.iTri = prev.iTri + prev.nTri;
                    cur.iVert = (ushort)(prev.iVert + prev.nVert);
                    if (prev.HasGeometry)
                    {
                        // Middle of strip -- back up and share geometry
                        cur.iTri -= kSharedTrisInSolidPair * NS;
                        cur.iVert -= (ushort)(kSharedVertsInSolidPair * NS);
                    }
                    cur.nTri = (ushort)(kTrisInSolid * NS);
                    cur.nVert = (ushort)(kVertsInSolid * NS);
                }

                m_knots[iKnot] = cur;
                prev = cur;
            }
        }

        void OnChanged_MakeVertsAndNormals(int iKnot0)
        {
            float belly = 1f / 8; // amount of belly rise per unit width
            // float theta = atan(belly/1);  belly angle
            float sinTheta, cosTheta;
            {
                float hypotenuse = Mathf.Sqrt(1 + belly * belly);
                sinTheta = belly / hypotenuse;
                cosTheta = 1 / hypotenuse;
            }


            // Invariant: there is a previous knot.
            Knot prev = m_knots[iKnot0 - 1];

            for (int iKnot = iKnot0; iKnot < m_knots.Count; ++iKnot)
            {
                // Invariant: all of prev's geometry (if any) is correct and up-to-date.
                // Thus, there is no need to modify anything shared with prev.
                Knot cur = m_knots[iKnot];

                if (cur.HasGeometry)
                {
                    bool isStart = !prev.HasGeometry;
                    bool isEnd = IsPenultimate(iKnot);

                    int j = 0;
                    // Top-Right
                    SetTri(cur.iTri, cur.iVert, j++, BRT, BMT, FMT);
                    SetTri(cur.iTri, cur.iVert, j++, BRT, FMT, FRT);
                    // Bottom-Right
                    SetTri(cur.iTri, cur.iVert, j++, BRB, FRB, FMB);
                    SetTri(cur.iTri, cur.iVert, j++, BRB, FMB, BMB);
                    // Bottom-Left
                    SetTri(cur.iTri, cur.iVert, j++, BMB, FLB, BLB);
                    SetTri(cur.iTri, cur.iVert, j++, BMB, FMB, FLB);
                    // Top-Left
                    SetTri(cur.iTri, cur.iVert, j++, BMT, FLT, FMT);
                    SetTri(cur.iTri, cur.iVert, j++, BMT, BLT, FLT);


                    Vector3 nUp = cur.nSurface;

                    if (isStart)
                    {
                        // Can't use prev.nRight, prev.nSurface; they're invalid if no geometry
                        float size = PressuredSize(prev.smoothedPressure);
                        float alpha = PressuredOpacity(prev.smoothedPressure);
                        Vector3 r = cur.nRight * (size / 2);
                        SetVert(cur.iVert, BRT, prev.point.m_Pos + r, nUp, m_Color, alpha);
                        SetVert(cur.iVert, BRB, prev.point.m_Pos + r, -nUp, m_Color, alpha);

                        SetVert(cur.iVert, BMT, prev.point.m_Pos, nUp, m_Color, alpha);
                        SetVert(cur.iVert, BMB, prev.point.m_Pos, -nUp, m_Color, alpha);

                        SetVert(cur.iVert, BLT, prev.point.m_Pos - r, nUp, m_Color, alpha);
                        SetVert(cur.iVert, BLB, prev.point.m_Pos - r, -nUp, m_Color, alpha);
                    }

                    {
                        float size = PressuredSize(cur.smoothedPressure);
                        float alpha = PressuredOpacity(cur.smoothedPressure);
                        Vector3 r = cur.nRight * (size / 2);
                        Vector3 u = cur.nSurface * (belly * size / 2);
                        Vector3 cosUp, sinRt;

                        if (isEnd)
                        {
                            // Pinch closed
                            cosUp = nUp;
                            u = sinRt = Vector3.zero;
                        }
                        else
                        {
                            cosUp = cosTheta * cur.nSurface;
                            sinRt = sinTheta * cur.nRight;
                        }

                        SetVert(cur.iVert, FRT, cur.point.m_Pos + r, cosUp + sinRt, m_Color, alpha);
                        SetVert(cur.iVert, FRB, cur.point.m_Pos + r, -cosUp + sinRt, m_Color, alpha);

                        SetVert(cur.iVert, FMT, cur.point.m_Pos + u, nUp, m_Color, alpha);
                        SetVert(cur.iVert, FMB, cur.point.m_Pos - u, -nUp, m_Color, alpha);

                        SetVert(cur.iVert, FLT, cur.point.m_Pos - r, cosUp - sinRt, m_Color, alpha);
                        SetVert(cur.iVert, FLB, cur.point.m_Pos - r, -cosUp - sinRt, m_Color, alpha);
                    }
                }

                prev = cur;
            }
        }

        bool IsPenultimate(int iKnot)
        {
            return (iKnot + 1 == m_knots.Count || !m_knots[iKnot + 1].HasGeometry);
        }

        void OnChanged_StretchUVs(int iChangedKnot)
        {
            // Back up knot to the start of the segment
            // Invariant: knot 0 never has geometry
            int iKnot0 = iChangedKnot;
            while (m_knots[iKnot0 - 1].HasGeometry)
            {
                iKnot0 -= 1;
            }

            while (iKnot0 < m_knots.Count)
            {
                // Update UVs for a single segment

                // Invariant: If iKnot0-1 exists, it has no geometry
                // IOW, iKnot0 is not the middle of a strip.
                // (We don't assume that it does have geometry, though)

                // Find length and end knot of segment
                float totalLength = 0;
                int iKnot1 = iKnot0;
                for (; iKnot1 < m_knots.Count; ++iKnot1)
                {
                    Knot cur = m_knots[iKnot1];
                    if (!cur.HasGeometry)
                    {
                        break;
                    }
                    totalLength += cur.length;
                }

                // Invariant: all knots in [iKnot0, iKnot1) have geometry
                // If iKnot1 exists, it does not have geometry

                float v0, v1;
                {
                    float random01 = m_rng.In01((int)m_knots[iKnot0].iVert - 1);
                    int numV = m_Desc.m_TextureAtlasV;
                    int iAtlas = (int)(random01 * 3331) % numV;
                    v0 = (iAtlas + m_TextureEdgeChop) / (float)numV;
                    v1 = (iAtlas + 1 - m_TextureEdgeChop) / (float)numV;
                }

                float distance = 0;
                for (int iKnot = iKnot0; iKnot < iKnot1; ++iKnot)
                {
                    Knot cur = m_knots[iKnot];
                    distance += cur.length;
                    float u = distance / totalLength;
                    if (iKnot == iKnot0)
                    {
                        SetUv0(cur.iVert, BLT, new Vector2(0, v0));
                        SetUv0(cur.iVert, BLB, new Vector2(0, v0));
                        SetUv0(cur.iVert, BRT, new Vector2(0, v1));
                        SetUv0(cur.iVert, BRB, new Vector2(0, v1));
                        SetUv0(cur.iVert, BMT, new Vector2(0, (v0 + v1) / 2));
                        SetUv0(cur.iVert, BMB, new Vector2(0, (v0 + v1) / 2));
                    }
                    SetUv0(cur.iVert, FLT, new Vector2(u, v0));
                    SetUv0(cur.iVert, FLB, new Vector2(u, v0));
                    SetUv0(cur.iVert, FRT, new Vector2(u, v1));
                    SetUv0(cur.iVert, FRB, new Vector2(u, v1));
                    SetUv0(cur.iVert, FMT, new Vector2(u, (v0 + v1) / 2));
                    SetUv0(cur.iVert, FMB, new Vector2(u, (v0 + v1) / 2));
                }

                iKnot0 = iKnot1 + 1;
            }
        }

        void OnChanged_DistanceUVs(int iKnot0)
        {
            Knot prev = m_knots[iKnot0 - 1];
            for (int iKnot = iKnot0; iKnot < m_knots.Count; ++iKnot)
            {
                Knot cur = m_knots[iKnot];

                if (cur.HasGeometry)
                {
                    // Textures are laid out so u goes along the strip,
                    // and v goes across the strip (from left to right)
                    float u0, v0, v1;
                    if (prev.HasGeometry)
                    {
                        // Middle of strip
                        Vector2 u0v0 = m_geometry.m_Texcoord0.v2[prev.iVert + FLT * NS];
                        u0 = u0v0.x;
                        v0 = u0v0.y;
                        v1 = m_geometry.m_Texcoord0.v2[prev.iVert + FRT * NS].y;
                    }
                    else
                    {
                        // Beginning of strip. Guaranteed no sharing, so previous
                        // vert should have a stable position.
                        float random01 = m_rng.In01(cur.iVert - 1);
                        u0 = random01;
                        int numV = m_Desc.m_TextureAtlasV;
                        int iAtlas = (int)(random01 * 3331) % numV;
                        v0 = (iAtlas + m_TextureEdgeChop) / (float)numV;
                        v1 = (iAtlas + 1 - m_TextureEdgeChop) / (float)numV;

                        SetUv0(cur.iVert, BLT, new Vector2(u0, v0));
                        SetUv0(cur.iVert, BLB, new Vector2(u0, v0));
                        SetUv0(cur.iVert, BRT, new Vector2(u0, v1));
                        SetUv0(cur.iVert, BRB, new Vector2(u0, v1));
                        SetUv0(cur.iVert, BMT, new Vector2(u0, (v0 + v1) / 2));
                        SetUv0(cur.iVert, BMB, new Vector2(u0, (v0 + v1) / 2));
                    }

                    float length = m_knots[iKnot].length;
                    float size = PressuredSize(m_knots[iKnot].smoothedPressure);
                    float u1 = u0 + m_Desc.m_TileRate * (length / size);
                    SetUv0(cur.iVert, FLT, new Vector2(u1, v0));
                    SetUv0(cur.iVert, FLB, new Vector2(u1, v0));
                    SetUv0(cur.iVert, FRT, new Vector2(u1, v1));
                    SetUv0(cur.iVert, FRB, new Vector2(u1, v1));
                    SetUv0(cur.iVert, FMT, new Vector2(u1, (v0 + v1) / 2));
                    SetUv0(cur.iVert, FMB, new Vector2(u1, (v0 + v1) / 2));
                }

                prev = cur;
            }
        }

        void OnChanged_Tangents(int iKnot0)
        {
            // TODO: just hardcode all this to fwd; no need to do all this junk

            Knot prev = m_knots[iKnot0 - 1];
            for (int iKnot = iKnot0; iKnot < m_knots.Count; ++iKnot)
            {
                Knot cur = m_knots[iKnot];
                if (cur.HasGeometry)
                {
                    // Omitting the *NS because this will never be double-sided
                    // Using the same vS for everything because I'm lazy
                    // Note that one side's normal map will have to be inverted because
                    // the seam shares verts (needs separate verts so they can flip w)
                    Vector3 vS, unused;
                    ComputeST(m_geometry.m_Vertices, m_geometry.m_Texcoord0.v2, cur.iVert, BRT, BMT, FMT,
                        out vS, out unused);

                    if (!prev.HasGeometry)
                    {
                        SetTangent(cur.iVert, BLT, vS);
                        SetTangent(cur.iVert, BLB, vS);
                        SetTangent(cur.iVert, BRT, vS);
                        SetTangent(cur.iVert, BRB, vS);
                        SetTangent(cur.iVert, BMT, vS);
                        SetTangent(cur.iVert, BMB, vS);
                    }

                    SetTangent(cur.iVert, FLT, vS);
                    SetTangent(cur.iVert, FLB, vS);
                    SetTangent(cur.iVert, FRT, vS);
                    SetTangent(cur.iVert, FRB, vS);
                    SetTangent(cur.iVert, FMT, vS);
                    SetTangent(cur.iVert, FMB, vS);
                }

                prev = cur;
            }
        }
    }
} // namespace TiltBrush
